package com.graph;



import java.util.*;

public class Solution863 {
    HashMap<TreeNode, TreeNode> map = new HashMap<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        HashMap<TreeNode, Boolean> visited = new HashMap<>();
        getParent(root);
        map.put(root, null);
        List<Integer> res = new ArrayList<>();
        LinkedList<TreeNode> queue = new LinkedList<>();
        queue.add(target);
        visited.put(target, true);
        while (k-- > 0) {
            int size = queue.size();
            while (size-- > 0) {
                TreeNode curNode = queue.removeFirst();
                if (curNode.left != null && !visited.getOrDefault(curNode.left, false)) {
                    queue.add(curNode.left);
                    visited.put(curNode.left, true);
                }
                if (curNode.right != null && !visited.getOrDefault(curNode.right, false)) {
                    queue.add(curNode.right);
                    visited.put(curNode.right, true);
                }
                if (map.get(curNode) != null && !visited.getOrDefault(map.get(curNode), false)) {
                    queue.add(map.get(curNode));
                    visited.put(map.get(curNode), true);
                }
            }
        }
        for (TreeNode t : queue) {
            res.add(t.val);
        }
        return res;
    }

    public void getParent(TreeNode root) {
        if (root.left != null) {
            map.put(root.left, root);
            getParent(root.left);
        }
        if (root.right != null) {
            map.put(root.right, root);
            getParent(root.right);
        }
    }
}

class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;

    TreeNode(int x) {
        val = x;
    }
}